package dns

import (
	"context"
	"fmt"
	"net"
	"sync/atomic"
	"time"

	"github.com/catroll/fastdns/internal/config"
	"github.com/catroll/fastdns/internal/logger"
)

// Server 表示 DNS 服务器
type Server struct {
	config     *config.Config
	resolver   *Resolver
	conn       *net.UDPConn
	queryCount *int64
	running    atomic.Bool
}

// NewServer 创建一个新的 DNS 服务器
func NewServer(cfg *config.Config, resolver *Resolver) *Server {
	return &Server{
		config:   cfg,
		resolver: resolver,
	}
}

// Start 启动 DNS 服务器
func (s *Server) Start() error {
	addr, err := net.ResolveUDPAddr("udp", s.config.Server.Listen)
	if err != nil {
		return fmt.Errorf("无法解析 DNS 服务器地址: %w", err)
	}

	conn, err := net.ListenUDP("udp", addr)
	if err != nil {
		return fmt.Errorf("无法启动 DNS 服务器: %w", err)
	}
	s.conn = conn

	logger.Info("DNS 服务器启动", "address", s.config.Server.Listen)
	s.running.Store(true)

	// 启动工作协程处理 DNS 请求
	go s.handleRequests()

	return nil
}

// Stop 停止 DNS 服务器
func (s *Server) Stop(ctx context.Context) error {
	if !s.running.Load() {
		return nil
	}

	logger.Info("DNS 服务器关闭中...")
	s.running.Store(false)

	if s.conn != nil {
		return s.conn.Close()
	}

	return nil
}

// handleRequests 处理传入的 DNS 请求
func (s *Server) handleRequests() {
	buffer := make([]byte, 512) // 标准 DNS 消息的最大 UDP 大小

	for s.running.Load() {
		n, addr, err := s.conn.ReadFromUDP(buffer)
		if err != nil {
			if !s.running.Load() {
				// 服务器正在关闭，这是预期的
				return
			}
			logger.Error("读取 DNS 请求失败", "error", err)
			continue
		}

		// 复制缓冲区以避免并发问题
		request := make([]byte, n)
		copy(request, buffer[:n])

		// 在新的 goroutine 中处理请求
		go s.handleDNSRequest(request, addr)
	}
}

// handleDNSRequest 处理单个 DNS 请求
func (s *Server) handleDNSRequest(request []byte, addr *net.UDPAddr) {
	// 解析 DNS 消息
	msg, err := parseDNSMessage(request)
	if err != nil {
		logger.Error("解析 DNS 消息失败", "error", err, "client", addr.String())
		return
	}

	logger.Debug("收到 DNS 查询",
		"client", addr.String(),
		"domain", msg.question.name,
		"type", msg.question.qtype)

	// 创建上下文，设置超时
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// 查询域名
	ips, err := s.resolver.LookupIP(ctx, msg.question.name)
	if err != nil {
		logger.Error("DNS 查询失败",
			"domain", msg.question.name,
			"error", err,
			"client", addr.String())

		// 发送错误响应
		response := createErrorResponse(msg)
		s.sendResponse(response, addr)
		return
	}

	// 创建响应
	response := createSuccessResponse(msg, ips, s.config.DNS.Cache.TTL)

	// 发送响应
	s.sendResponse(response, addr)

	logger.Debug("发送 DNS 响应",
		"client", addr.String(),
		"domain", msg.question.name,
		"ips_count", len(ips))

	// 增加查询计数
	if s.queryCount != nil {
		atomic.AddInt64(s.queryCount, 1)
	}
}

// sendResponse 发送 DNS 响应
func (s *Server) sendResponse(response []byte, addr *net.UDPAddr) {
	if _, err := s.conn.WriteToUDP(response, addr); err != nil {
		logger.Error("发送 DNS 响应失败", "error", err, "client", addr.String())
	}
}

// SetQueryCounter 设置查询计数器引用
func (s *Server) SetQueryCounter(counter *int64) {
	s.queryCount = counter
}

// DNSMessage 表示 DNS 消息
type DNSMessage struct {
	id       uint16
	flags    uint16
	question DNSQuestion
}

// DNSQuestion 表示 DNS 问题部分
type DNSQuestion struct {
	name  string
	qtype uint16
	class uint16
}

// 常量定义
const (
	// DNS 记录类型
	TypeA     uint16 = 1  // IPv4 地址
	TypeAAAA  uint16 = 28 // IPv6 地址
	TypeCNAME uint16 = 5  // 规范名称
	TypeMX    uint16 = 15 // 邮件交换
	TypeTXT   uint16 = 16 // 文本记录

	// DNS 响应码
	RcodeNoError        uint16 = 0 // 没有错误
	RcodeFormatError    uint16 = 1 // 格式错误
	RcodeServerFailure  uint16 = 2 // 服务器失败
	RcodeNameError      uint16 = 3 // 域名不存在
	RcodeNotImplemented uint16 = 4 // 查询类型不支持
	RcodeRefused        uint16 = 5 // 查询被拒绝

	// DNS 标志
	FlagQR uint16 = 0x8000 // 查询/响应标志
	FlagAA uint16 = 0x0400 // 权威应答标志
	FlagTC uint16 = 0x0200 // 截断标志
	FlagRD uint16 = 0x0100 // 递归查询标志
	FlagRA uint16 = 0x0080 // 递归可用标志
)

// parseDNSMessage 解析 DNS 消息
func parseDNSMessage(data []byte) (*DNSMessage, error) {
	if len(data) < 12 {
		return nil, fmt.Errorf("DNS 消息太短")
	}

	// 解析 DNS 头部
	id := uint16(data[0])<<8 | uint16(data[1])
	flags := uint16(data[2])<<8 | uint16(data[3])
	qdcount := uint16(data[4])<<8 | uint16(data[5])

	if qdcount != 1 {
		return nil, fmt.Errorf("仅支持单个问题")
	}

	// 解析问题部分
	offset := 12
	name, newOffset, err := parseDomainName(data, offset)
	if err != nil {
		return nil, err
	}
	offset = newOffset

	if offset+4 > len(data) {
		return nil, fmt.Errorf("DNS 消息太短，无法解析问题类型和类")
	}

	qtype := uint16(data[offset])<<8 | uint16(data[offset+1])
	class := uint16(data[offset+2])<<8 | uint16(data[offset+3])

	return &DNSMessage{
		id:    id,
		flags: flags,
		question: DNSQuestion{
			name:  name,
			qtype: qtype,
			class: class,
		},
	}, nil
}

// parseDomainName 解析 DNS 域名
func parseDomainName(data []byte, offset int) (string, int, error) {
	if offset >= len(data) {
		return "", 0, fmt.Errorf("偏移量超出范围")
	}

	var name string
	var jumped bool
	var jumpsPerformed int
	var currentOffset = offset
	var nextByte byte

	for {
		if currentOffset >= len(data) {
			return "", 0, fmt.Errorf("偏移量超出范围")
		}

		nextByte = data[currentOffset]
		currentOffset++

		if nextByte == 0 {
			break
		}

		// 处理压缩指针
		if (nextByte & 0xC0) == 0xC0 {
			if currentOffset >= len(data) {
				return "", 0, fmt.Errorf("偏移量超出范围")
			}

			if !jumped {
				offset = currentOffset + 1
			}

			currentOffset = int(uint16(nextByte&0x3F)<<8 | uint16(data[currentOffset]))
			jumped = true
			jumpsPerformed++

			if jumpsPerformed > 10 {
				return "", 0, fmt.Errorf("过多的指针跳转")
			}

			continue
		}

		// 读取标签
		labelLen := int(nextByte)
		if currentOffset+labelLen > len(data) {
			return "", 0, fmt.Errorf("标签长度超出范围")
		}

		if name != "" {
			name += "."
		}

		name += string(data[currentOffset : currentOffset+labelLen])
		currentOffset += labelLen
	}

	if !jumped {
		offset = currentOffset
	}

	return name, offset, nil
}

// createSuccessResponse 创建成功的 DNS 响应
func createSuccessResponse(req *DNSMessage, ips []net.IP, ttl time.Duration) []byte {
	// 创建响应头部
	response := make([]byte, 512)

	// ID (2 字节)
	response[0] = byte(req.id >> 8)
	response[1] = byte(req.id)

	// 标志 (2 字节): QR=1, RD=1, RA=1, RCODE=0
	flags := FlagQR | (req.flags & FlagRD) | FlagRA
	response[2] = byte(flags >> 8)
	response[3] = byte(flags)

	// QDCOUNT (2 字节): 问题数量 = 1
	response[4] = 0
	response[5] = 1

	// ANCOUNT (2 字节): 回答数量
	ancount := uint16(len(ips))
	response[6] = byte(ancount >> 8)
	response[7] = byte(ancount)

	// NSCOUNT (2 字节): 权威名称服务器数量 = 0
	response[8] = 0
	response[9] = 0

	// ARCOUNT (2 字节): 附加记录数量 = 0
	response[10] = 0
	response[11] = 0

	// 复制问题部分
	offset := 12

	// 写入域名
	offset = writeDomainName(response, offset, req.question.name)

	// 写入问题类型和类
	response[offset] = byte(req.question.qtype >> 8)
	response[offset+1] = byte(req.question.qtype)
	response[offset+2] = byte(req.question.class >> 8)
	response[offset+3] = byte(req.question.class)
	offset += 4

	// 写入回答部分
	ttlSeconds := uint32(ttl.Seconds())

	for _, ip := range ips {
		// 写入域名指针
		response[offset] = 0xC0
		response[offset+1] = 12
		offset += 2

		// 写入记录类型
		var recordType uint16
		if ip.To4() != nil {
			recordType = TypeA
		} else {
			recordType = TypeAAAA
		}
		response[offset] = byte(recordType >> 8)
		response[offset+1] = byte(recordType)
		offset += 2

		// 写入记录类
		response[offset] = byte(req.question.class >> 8)
		response[offset+1] = byte(req.question.class)
		offset += 2

		// 写入 TTL
		response[offset] = byte(ttlSeconds >> 24)
		response[offset+1] = byte(ttlSeconds >> 16)
		response[offset+2] = byte(ttlSeconds >> 8)
		response[offset+3] = byte(ttlSeconds)
		offset += 4

		// 写入数据长度和数据
		if ip.To4() != nil {
			// IPv4 地址 (4 字节)
			response[offset] = 0
			response[offset+1] = 4
			offset += 2

			ipv4 := ip.To4()
			copy(response[offset:], ipv4)
			offset += 4
		} else {
			// IPv6 地址 (16 字节)
			response[offset] = 0
			response[offset+1] = 16
			offset += 2

			ipv6 := ip.To16()
			copy(response[offset:], ipv6)
			offset += 16
		}
	}

	return response[:offset]
}

// createErrorResponse 创建错误的 DNS 响应
func createErrorResponse(req *DNSMessage) []byte {
	// 创建响应头部
	response := make([]byte, 512)

	// ID (2 字节)
	response[0] = byte(req.id >> 8)
	response[1] = byte(req.id)

	// 标志 (2 字节): QR=1, RD=1, RA=1, RCODE=2 (服务器失败)
	flags := FlagQR | (req.flags & FlagRD) | FlagRA | RcodeServerFailure
	response[2] = byte(flags >> 8)
	response[3] = byte(flags)

	// QDCOUNT (2 字节): 问题数量 = 1
	response[4] = 0
	response[5] = 1

	// ANCOUNT (2 字节): 回答数量 = 0
	response[6] = 0
	response[7] = 0

	// NSCOUNT (2 字节): 权威名称服务器数量 = 0
	response[8] = 0
	response[9] = 0

	// ARCOUNT (2 字节): 附加记录数量 = 0
	response[10] = 0
	response[11] = 0

	// 复制问题部分
	offset := 12

	// 写入域名
	offset = writeDomainName(response, offset, req.question.name)

	// 写入问题类型和类
	response[offset] = byte(req.question.qtype >> 8)
	response[offset+1] = byte(req.question.qtype)
	response[offset+2] = byte(req.question.class >> 8)
	response[offset+3] = byte(req.question.class)
	offset += 4

	return response[:offset]
}

// writeDomainName 将域名写入 DNS 消息
func writeDomainName(buffer []byte, offset int, domain string) int {
	if domain == "" {
		buffer[offset] = 0
		return offset + 1
	}

	var start, end int
	for end <= len(domain) {
		if end == len(domain) || domain[end] == '.' {
			labelLen := end - start
			buffer[offset] = byte(labelLen)
			offset++

			copy(buffer[offset:], domain[start:end])
			offset += labelLen

			start = end + 1
		}
		end++
	}

	buffer[offset] = 0 // 域名结束标记
	return offset + 1
}
